{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from charm.toolbox.pairinggroup import PairingGroup, G1, GT, ZR, pair as e\n",
    "from node import *\n",
    "from lsssutil import *\n",
    "from Util import a1_param_generator\n",
    "import numpy as np\n",
    "from functools import reduce\n",
    "\n",
    "a1_conf_path = '../pbc.conf.d/a1.properties'\n",
    "primes = a1_param_generator(60, 3, a1_conf_path)\n",
    "group = PairingGroup(a1_conf_path, param_file = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![example](../imgs/cp_abe_example.jpg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## setup\n",
    "![setup](./article_imgs/1.jpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def setup(sec_lambda, U):\n",
    "    alpha, a = group.random(ZR, 2)\n",
    "    g = group.random(G1) ** (group.order() // primes[0])\n",
    "    T = [None for _ in range(len(U))]\n",
    "    for i in U:\n",
    "        T[i] = g ** group.random(ZR)\n",
    "    \n",
    "    return {'PK': {'N': group.order(), 'g': g, 'g_a': g ** a, 'egg_alpha': e(g, g) ** alpha, 'T': T},\n",
    "            'MSK': {'alpha': alpha, 'X3': group.random(G1) ** (group.order() // primes[2])}}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## key generator\n",
    "![key generator](./article_imgs/2.jpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def key_generate(MSK, S, PK):\n",
    "    t = group.random(ZR)\n",
    "    R0, R0_minute = group.random(G1) ** (group.order() // primes[2]), group.random(G1) ** (group.order() // primes[2])\n",
    "    g, g_a, alpha, T = PK['g'], PK['g_a'], MSK['alpha'], PK['T']\n",
    "    K = (g ** alpha) * (g_a ** t) * R0\n",
    "    L = (g ** t) * R0_minute\n",
    "    Ki = [(T[i] ** t) * (group.random(G1) ** (group.order() // primes[2])) for i in S]\n",
    "    \n",
    "    return {'S': S, 'K': K, 'L': L, 'Ki': Ki}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## encryption\n",
    "![encryption](./article_imgs/3.jpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def encryption(A, PK, M):\n",
    "    l, n = A.shape\n",
    "    g, g_a, egg_alpha, T = PK['g'], PK['g_a'], PK['egg_alpha'], PK['T']\n",
    "    s = group.random(ZR)\n",
    "    v = [s] + [group.random(ZR) for _ in range(n - 1)]\n",
    "    C, C_minute = M * (egg_alpha ** s), g ** s\n",
    "    Cx, Dx = [None] * l, [None] * l\n",
    "    for i in range(l):\n",
    "        rx = group.random(ZR)\n",
    "        Cx[i] = (g_a ** (A[i].dot(v))) * (T[i] ** (-rx))\n",
    "        Dx[i] = g ** rx\n",
    "\n",
    "    return {'C': C, 'C_minute': C_minute, 'Cx': Cx, 'Dx': Dx}\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## decrypt\n",
    "![decrypt](./article_imgs/4.jpg)\n",
    "\\\n",
    "正确性证明\n",
    "\\begin{align*} \n",
    "  \\prod_{\\rho \\left( x \\right) \\in S}{\\left( e\\left( C_x,L \\right) e\\left( D_x,K_{\\rho (x)} \\right) \\right) ^{w_x}} &= \\prod_{\\rho (x)\\in S}{\\left( e\\left( g^{aA_x\\cdot v}T_{\\rho (x)}^{-r_x}, g^tR_{0}^{'} \\right) e\\left( g^{r_x},T_{i}^{t}R_i \\right) \\right) ^{w_x}}\n",
    "\\\\ &= \\prod_{\\rho (x)\\in S}{\\left( e\\left( g^{aA_x\\cdot v}T_{\\rho (x)}^{-r_x},g^t \\right) e\\left( g^{r_x},T_{i}^{t} \\right) \\right) ^{w_x}}\n",
    "\\\\ &= \\prod_{\\rho (x)\\in S}{\\left( e\\left( g^{aA_x\\cdot v},g^t \\right) e\\left( T_{\\rho \\left( x \\right)}^{-r_x},g^t \\right) e\\left( g^t,T_{i}^{r_x} \\right) \\right) ^{w_x}}\n",
    "\\\\ &= \\prod_{\\rho (x)\\in S}{\\left( e\\left( g^{aA_x\\cdot v},g^t \\right) e\\left( g^t,T_{i}^{0} = \\mathrm{O} \\right) \\right) ^{w_x}}\n",
    "\\\\ &= \\prod_{\\rho \\left( x \\right) \\in S}{e\\left( g^{aA_x\\cdot v},g^t \\right) ^{w_x}}\n",
    "\\\\ &= e\\left( g,g \\right) ^{at\\sum{\\left( A_x\\cdot v \\right) w_x}}\n",
    "\\\\ &= e\\left( g,g \\right) ^{at\\sum{\\lambda _xw_x}}\n",
    "\\\\ &= e\\left( g,g \\right) ^{ats}\n",
    "\\end{align*}\n",
    "于是\n",
    "$$\n",
    "e\\left( C',K \\right) =e\\left( g^s,g^{\\alpha}g^{at} \\right) =e\\left( g^s,g^{\\alpha} \\right) e\\left( g^s,g^{at} \\right) =e\\left( g,g \\right) ^{\\alpha s}e\\left( g,g \\right) ^{ats}\n",
    "$$\n",
    "$$\n",
    "M=C/e\\left( g,g \\right) ^{\\alpha s}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def decrypt(A, CT, PK, SK):\n",
    "    C, C_minute, Cx, Dx = CT['C'], CT['C_minute'], CT['Cx'], CT['Dx']\n",
    "    S, K, L, Ki = SK['S'], SK['K'], SK['L'], SK['Ki']\n",
    "    Ax = A[S,:].T\n",
    "    row, _ = Ax.shape\n",
    "    b = np.array([1] + [0 for _ in range(row - 1)], dtype = int)\n",
    "    wx = np.rint(np.linalg.lstsq(Ax, b, rcond = None)[0]).astype(int).tolist()\n",
    "    e_wx = []\n",
    "    for index, attr in enumerate(S):\n",
    "        e_wx.append((e(Cx[attr], L) * e(Dx[attr], Ki[index])) ** group.init(ZR, wx[index]))\n",
    "    prod = reduce(lambda e1, e2: e1 * e2, e_wx)\n",
    "    egg_alpha_s = e(C_minute, K) / prod\n",
    "    \n",
    "    return C / egg_alpha_s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## test\n",
    "需要注意的是，不能同时传入或节点的两个子节点，因为这时的$A_xw_x = \\left(1, 0, ..., 0\\right)$是一个不定方程，numpy得不到整数类型的特解，也就不能进行双线性运算。解密不成功。\\\n",
    "比如下面的S = [0, 3, 4]本来可以正常解密，结果解密不正确，正确的输入应该是S = [0, 3] 或 S = [0, 4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1099911053833442339924112560261511518064759114939708232, 52203878605640074084434760195387730592499702629859205]\n",
      "LSSS矩阵：\n",
      "[[ 0 -1  0]\n",
      " [ 0  0 -1]\n",
      " [ 1  1  1]\n",
      " [ 1  1  0]\n",
      " [ 1  1  0]]\n",
      "解密成功，解密结果：\n",
      "[1099911053833442339924112560261511518064759114939708232, 52203878605640074084434760195387730592499702629859205]\n"
     ]
    }
   ],
   "source": [
    "U = range(5)\n",
    "S = [3, 0, 1, 2]\n",
    "M = group.random(GT)\n",
    "print(M)\n",
    "\n",
    "# 布尔访问树结构\n",
    "root = Node.threshold_node(Gate(2, 2), [\n",
    "    Node.attr_node('A'),\n",
    "    Node.threshold_node(Gate(1, 2), [\n",
    "        Node.threshold_node(Gate(2, 2), [Node.attr_node('B'), Node.attr_node('C')]),\n",
    "        Node.threshold_node(Gate(1, 2), [Node.attr_node('D'), Node.attr_node('E')])\n",
    "    ])\n",
    "])\n",
    "\n",
    "# 获取布尔访问树的LSSS矩阵，并转为numpy array\n",
    "matrix = lewko_waters(root)\n",
    "npmatrix = np.array(matrix, dtype = int)\n",
    "print(f'LSSS矩阵：\\n{npmatrix}')\n",
    "\n",
    "# 生成公共参数\n",
    "params = setup(0, U)\n",
    "\n",
    "# 获取私钥\n",
    "SK = key_generate(params['MSK'], S, params['PK'])\n",
    "\n",
    "# 加密消息M，返回密文\n",
    "CT = encryption(npmatrix, params['PK'], M)\n",
    "\n",
    "# 解密消息\n",
    "dec_M = decrypt(npmatrix, CT, params['PK'], SK)\n",
    "\n",
    "assert dec_M == M, '解密失败'\n",
    "print(f'解密成功，解密结果：\\n{dec_M}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.9 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "e5030792b3492f6b12d94f1f48beca3d8e59ec05fd59d0aaaa48e684281ed297"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
